# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import logging
import gzip

from collections import namedtuple

from coin.base.param_util import to_list
from coin.base.datetime_util import to_datetime
from coin.strategy.marketsim.core import spec_pb2


def open_file(filepath, *args, **kwargs):
  if filepath.endswith('.gz'):
    return gzip.open(filepath, *args, **kwargs)
  else:
    return open(filepath, *args, **kwargs)


def read_sim_result(filepaths):
  results = {}
  for filepath in sorted(to_list(filepaths)):
    with open_file(filepath, 'rb') as f:
      res = spec_pb2.MarketSimResultProto()
      res.ParseFromString(f.read())
      results[filepath] = res
  return results


Fill = namedtuple('Fill', ['ts', 'price', 'signed_qty'])


def collect_fills(*, sim_results=None, filepaths=None, filter_func=None, verbose=True):
  assert (sim_results is None) != (filepaths is None)
  filter_func = filter_func or (lambda _: True)
  if sim_results is None:
    sim_results = read_sim_result(filepaths)

  fills = {}
  last_ts = 0
  for _, res in sim_results.items():
    for execution in res.executions:
      if not (execution.exec_type == spec_pb2.EXEC_TYPE_FILLED and filter_func(execution)):
        continue
      if execution.symbol not in fills:
        fills[execution.symbol] = []

      fill_list = fills[execution.symbol]
      exec_ts = execution.exec_ts + 1  # Shift 1ns
      if fill_list and exec_ts <= last_ts:
        exec_ts = last_ts + 1
      last_ts = exec_ts

      fill = Fill(exec_ts, execution.fill_price, execution.fill_signed_qty)
      fill_list.append(fill)

  for k in list(fills.keys()):
    fills[k] = sorted(fills[k], key=lambda f: f.ts)

  if verbose:
    from_ts, to_ts = 2**64, -(2**64)
    for symbol, fills_per_symbol in fills.items():
      from_ts = min(fills_per_symbol[0].ts, from_ts)
      to_ts = max(fills_per_symbol[-1].ts, to_ts)
      logging.info('%s: %s', symbol, len(fills_per_symbol))
    logging.info('Range: %s - %s' % (to_datetime(from_ts), to_datetime(to_ts)))

  return fills
